import gradio as gr
import joblib
import numpy as np
from PIL import Image
import os

# 配置模型文件路径
MODEL_PATH = 'best_knn_model.pkl'

# 尝试加载保存的KNN模型
try:
    knn_model = joblib.load(MODEL_PATH)
    print(f"Model loaded successfully from {MODEL_PATH}")
except Exception as e:
    print(f"Failed to load model: {e}")
    raise

# 图像预处理函数
def preprocess_image(image):
    try:
        # 确保输入是有效的
        if not isinstance(image, dict) or 'layers' not in image:
            raise ValueError("Invalid image format")

        # 提取 alpha 通道并调整大小为 8x8
        alpha_channel = image['layers'][0][:, :, 3]
        img_pil = Image.fromarray(alpha_channel).resize((8, 8))

        # 将图像转换为 numpy 数组，并归一化为 0-16 范围
        img_array = (np.array(img_pil).reshape(1, -1) / 255.0) * 16

        # 使用 KNN 模型进行预测并返回结果
        return int(knn_model.predict(img_array)[0])
    except Exception as e:
        print(f"Error processing image: {e}")
        return "Error processing image"

# 创建 Gradio 界面
iface = gr.Interface(
    fn=preprocess_image,  # 设置要调用的预测函数
    inputs=gr.Sketchpad(canvas_size=(300, 300), label="绘制数字"),  # 设置输入类型为手写板，并调整画板大小为300x300
    outputs=gr.Label(num_top_classes=1),  # 设置输出类型为标签，显示预测结果
    title="手写数字识别",  # 设置界面的标题
    description="请在画布上绘制一个0-9之间的数字，点击提交预测。"  # 设置界面的描述
)

# 启动 Gradio 界面
iface.launch(share=True)